{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyMOm3tUW2I1qdZ0w0WEW6+M",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LC1332/Chat-Haruhi-Suzumiya/blob/main/notebook/%E8%B6%85%E8%B6%8A%E5%AE%A1%E6%A0%B8%E8%BE%B9%E7%95%8C%E7%9A%84embedding%E6%8A%BD%E5%8F%96%E5%92%8Cjsonl%E7%94%9F%E6%88%90.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "对几个超越审核边界的人物进行抽取\n",
        "\n",
        "implemented by 李鲁鲁\n",
        "\n",
        "- [x] 对接一下bge_zh_small\n",
        "- [ ] 解析单个的审核文本\n",
        "- [ ] 对text字段进行base编码\n",
        "- [ ] 重写一个外部的package role函数\n",
        "- [ ] (opt*) 抽取bge_small的同时抽取luotuo\n",
        "\n",
        "原版程序在\n",
        "\n",
        "https://github.com/LC1332/Chat-Haruhi-Suzumiya/blob/main/notebook/%E5%B0%9D%E8%AF%95%E8%BD%AC%E5%8C%96RoleLLM%E4%B8%BAHaruhi%E6%A0%BC%E5%BC%8F.ipynb\n",
        "\n",
        "和\n",
        "\n",
        "\n",
        "https://github.com/LC1332/Chat-Haruhi-Suzumiya/blob/main/characters/novel_collecting/chatbot%E7%9A%84embedding%E6%8A%BD%E5%8F%96%E5%92%8Cjsonl%E7%94%9F%E6%88%90.ipynb\n",
        "\n"
      ],
      "metadata": {
        "id": "pTr24hXRL2rh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "先解析文本"
      ],
      "metadata": {
        "id": "uYdJDq0WRk2o"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "建立txt specific file-role name的映射"
      ],
      "metadata": {
        "id": "GTXOeHKpRsqX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "specific_fname2role = {\"艾希2\":\"艾希\"}"
      ],
      "metadata": {
        "id": "UKG2ERVGRmZo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "input_folder = \"/content/input/\"\n",
        "\n",
        "texts = []\n",
        "\n",
        "import os\n",
        "\n",
        "save_names = []\n",
        "\n",
        "#read txt files from folder\n",
        "for fname in os.listdir(input_folder):\n",
        "    if fname.endswith(\".txt\"):\n",
        "        with open(os.path.join(input_folder, fname), 'r', encoding='utf-8') as f:\n",
        "            texts.append(f.read())\n",
        "            save_names.append(fname[:-4])"
      ],
      "metadata": {
        "id": "D0gmjFAMShoW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(fname[:-4])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hlMuYodwb2MD",
        "outputId": "d9f32aa3-6891-4c00-fd94-12ce3d0d4d82"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "安莉娜\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(texts[0][:150])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yAJoE0lxTHMn",
        "outputId": "e573359a-faa2-4278-89f6-6d6ce819b979"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "I want you to act like 女贤者 from 情色小说 Erotica.\n",
            "Please try to answer with similar language style from the Erotica.\n",
            "\n",
            "女贤者没有性相关的经验，谈论相关的事情会觉得很羞耻\n",
            "女贤者是个洁白无瑕的\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "role_datas = []\n",
        "\n",
        "for text in texts:\n",
        "    chunks = text.split(\"###\")\n",
        "    for chunk in chunks:\n",
        "        if \"###\" in chunk:\n",
        "            print(chunk)\n",
        "    role_data = {\n",
        "        \"system_prompt\":chunks[0].strip(),\n",
        "        \"texts\":chunks[1:]\n",
        "    }\n",
        "    role_datas.append(role_data)"
      ],
      "metadata": {
        "id": "s5OUmKk4TPnl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(save_names)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nEuR1No1b7jp",
        "outputId": "fcb9483e-0298-4fbb-c623-963c0d531c95"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['贤者', '艾希2', '艾希', '安莉娜']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content\n",
        "!rm -rf /content/Haruhi-2-Dev\n",
        "!git clone https://github.com/LC1332/Haruhi-2-Dev\n",
        "%cd /content/Haruhi-2-Dev\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0IASe5tHhu_y",
        "outputId": "a2b029a2-f5dc-4ae7-aef4-5e142a530da1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content\n",
            "Cloning into 'Haruhi-2-Dev'...\n",
            "remote: Enumerating objects: 929, done.\u001b[K\n",
            "remote: Counting objects: 100% (236/236), done.\u001b[K\n",
            "remote: Compressing objects: 100% (127/127), done.\u001b[K\n",
            "remote: Total 929 (delta 159), reused 162 (delta 109), pack-reused 693\u001b[K\n",
            "Receiving objects: 100% (929/929), 106.26 MiB | 14.89 MiB/s, done.\n",
            "Resolving deltas: 100% (489/489), done.\n",
            "/content/Haruhi-2-Dev\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
      ],
      "metadata": {
        "id": "zAGJEBi8r8Hc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(device)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6Fra2xvfsLeD",
        "outputId": "0fcc4827-780e-4cda-a846-884f03c3bbc5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "cuda\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip -q install openai tiktoken langchain chromadb"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pKm2xA5UUtB-",
        "outputId": "8b667f1b-b2ae-454e-816e-6d20d9515dd8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m221.4/221.4 kB\u001b[0m \u001b[31m4.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m38.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m83.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m502.4/502.4 kB\u001b[0m \u001b[31m47.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.0/75.0 kB\u001b[0m \u001b[31m9.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m181.5/181.5 kB\u001b[0m \u001b[31m22.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m48.2/48.2 kB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m94.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m92.9/92.9 kB\u001b[0m \u001b[31m13.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.7/59.7 kB\u001b[0m \u001b[31m8.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.4/5.4 MB\u001b[0m \u001b[31m110.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.4/6.4 MB\u001b[0m \u001b[31m119.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m57.9/57.9 kB\u001b[0m \u001b[31m3.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m105.3/105.3 kB\u001b[0m \u001b[31m14.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.3/67.3 kB\u001b[0m \u001b[31m9.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m699.4/699.4 kB\u001b[0m \u001b[31m47.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m57.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m72.6/72.6 kB\u001b[0m \u001b[31m8.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.4/49.4 kB\u001b[0m \u001b[31m5.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.0/67.0 kB\u001b[0m \u001b[31m8.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.9/76.9 kB\u001b[0m \u001b[31m10.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m8.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m143.8/143.8 kB\u001b[0m \u001b[31m18.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.0/46.0 kB\u001b[0m \u001b[31m6.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m50.8/50.8 kB\u001b[0m \u001b[31m7.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m341.4/341.4 kB\u001b[0m \u001b[31m38.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.4/3.4 MB\u001b[0m \u001b[31m103.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m72.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m130.2/130.2 kB\u001b[0m \u001b[31m15.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.8/86.8 kB\u001b[0m \u001b[31m11.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for pypika (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "lida 0.0.10 requires kaleido, which is not installed.\n",
            "lida 0.0.10 requires python-multipart, which is not installed.\n",
            "llmx 0.0.15a0 requires cohere, which is not installed.\n",
            "tensorflow-probability 0.22.0 requires typing-extensions<4.6.0, but you have typing-extensions 4.8.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from ChatHaruhi.utils import float_array_to_base64, base64_to_float_array"
      ],
      "metadata": {
        "id": "7RJzZDcailY_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "这里我们希望用两种embeddings去写"
      ],
      "metadata": {
        "id": "xC1xzTw-U-mW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from ChatHaruhi.utils import get_bge_zh_embedding, luotuo_openai_embedding\n",
        "# for english character should use get_bge_en_embedding instead"
      ],
      "metadata": {
        "id": "s00HKcoLVSCV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "embed_bge = get_bge_zh_embedding(\"测试\")\n",
        "print(embed_bge[:5])\n",
        "print(len(embed_bge))"
      ],
      "metadata": {
        "id": "jNtC1leIOpYl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "embed_luotuo = luotuo_openai_embedding(\"测试\")\n",
        "print(embed_luotuo[:5])\n",
        "print(len(embed_luotuo))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GLDTxdPMVkEh",
        "outputId": "a34a88da-3838-4ef8-d8fc-e7e783f4a2c4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[-0.011614263989031315, -0.007930075749754906, 0.01195011381059885, -0.011979658156633377, -0.03503050282597542]\n",
            "1536\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "embed_funcs = [luotuo_openai_embedding, get_bge_zh_embedding]\n",
        "column_names = [\"luotuo_openai\", \"bge_zh_s15\"]"
      ],
      "metadata": {
        "id": "5Ydo5mXSVuiz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "TODO: 把这两个移动到util里面去"
      ],
      "metadata": {
        "id": "2njCH_y4ah6L"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import base64\n",
        "import struct\n",
        "\n",
        "def string_to_base64(text):\n",
        "    byte_array = b''\n",
        "    for char in text:\n",
        "        num_bytes = char.encode('utf-8')\n",
        "        byte_array += num_bytes\n",
        "\n",
        "    base64_data = base64.b64encode(byte_array)\n",
        "    return base64_data.decode('utf-8')\n",
        "\n",
        "def base64_to_string(base64_data):\n",
        "    byte_array = base64.b64decode(base64_data)\n",
        "    text = byte_array.decode('utf-8')\n",
        "    return text"
      ],
      "metadata": {
        "id": "Cd5wLn8lagbF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "from tqdm import tqdm\n",
        "def package_role( system_prompt, texts, embed_funcs, embed_names , if_enc = False):\n",
        "    datas = []\n",
        "    if True:\n",
        "        data = {}\n",
        "        if if_enc:\n",
        "            data['enc_text'] = string_to_base64(system_prompt)\n",
        "        else:\n",
        "            data['text'] = system_prompt\n",
        "\n",
        "        for embed_name in embed_names:\n",
        "            data[embed_name] = 'system_prompt'\n",
        "        datas.append(data)\n",
        "\n",
        "    if True:\n",
        "        data = {}\n",
        "\n",
        "        if if_enc:\n",
        "            data['enc_text'] = string_to_base64(\"Reserve Config Setting Here\")\n",
        "        else:\n",
        "            data['text'] = \"Reserve Config Setting Here\"\n",
        "\n",
        "        for embed_name in embed_names:\n",
        "            data[embed_name] = 'config'\n",
        "        datas.append(data)\n",
        "\n",
        "    for text in tqdm(texts):\n",
        "        data = {}\n",
        "        text = text.strip()\n",
        "        if if_enc:\n",
        "            data['enc_text'] = string_to_base64(text)\n",
        "        else:\n",
        "            data['text'] = text\n",
        "        for embed_func, embed_name in zip(embed_funcs, embed_names):\n",
        "            current_vec = embed_func( text )\n",
        "            encode_vec = float_array_to_base64(current_vec)\n",
        "            data[embed_name] = encode_vec\n",
        "        datas.append(data)\n",
        "    return datas"
      ],
      "metadata": {
        "id": "AMHjBKwVVGkQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(role_datas[0].keys())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MZBPEm8PZTsD",
        "outputId": "aed6e3d2-8226-4b00-bec9-aa163a568bff"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dict_keys(['system_prompt', 'texts'])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "packaged_datas = package_role(\n",
        "    role_datas[0]['system_prompt'],\n",
        "    role_datas[0]['texts'][:2],\n",
        "    embed_funcs,\n",
        "    column_names,\n",
        "    if_enc = True\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JQhGhP62ZQIU",
        "outputId": "d8ad8729-d01a-4628-876d-e085ccfb8bac"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 2/2 [00:00<00:00,  2.10it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "import json\n",
        "\n",
        "def write_jsonl(datas, file_name):\n",
        "    with open(file_name, 'w', encoding='utf-8') as f:\n",
        "        for data in datas:\n",
        "            json_str = json.dumps(data, ensure_ascii=False)\n",
        "            f.write(json_str+\"\\n\")\n",
        "\n",
        "for data, save_name in zip(role_datas,save_names):\n",
        "    save_path = \"/content/output/\" + save_name + \".jsonl\"\n",
        "    packaged_data = package_role(\n",
        "        data['system_prompt'],\n",
        "        data['texts'],\n",
        "        embed_funcs,\n",
        "        column_names,\n",
        "        if_enc=True\n",
        "    )\n",
        "    write_jsonl(packaged_data, save_path)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hvEuX289cAKK",
        "outputId": "0b4bd81d-6d4f-4436-9391-14f4e72504f4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 35/35 [00:16<00:00,  2.11it/s]\n",
            "100%|██████████| 14/14 [00:06<00:00,  2.28it/s]\n",
            "100%|██████████| 17/17 [00:07<00:00,  2.33it/s]\n",
            "100%|██████████| 24/24 [00:10<00:00,  2.39it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "qW91sq8uwhyE"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}